import cv2 as cv
import pickle
import os
import numpy as np


win_size = (128, 128)
block_size = (16, 16)
block_stride = (8, 8)
cell_size = (8, 8)
n_bin = 9
hog = cv.HOGDescriptor(win_size, block_size, block_stride, cell_size, n_bin)


def get_feature(path):
    img = cv.imread(path)
    img = cv.resize(img, (128, 128))
    feature = hog.compute(img=img, winStride=(8, 8))
    feature = np.array(feature)
    feature = feature.reshape(-1)
    return feature


def main():
    data_dir = 'train/'
    contents = os.listdir(data_dir)
    labels = [content.split('.')[0] for content in contents]
    labels = [[1] if label == 'cat' else [0] for label in labels]
    labels = np.array(labels)
    features = np.array([get_feature(os.path.join(data_dir, content)) for content in contents])

    print(labels.shape)
    print(features.shape)
    pickle.dump(features, open('data/features.pkl', 'wb'))
    pickle.dump(labels, open('data/labels.pkl', 'wb'))


if __name__ == '__main__':
    main()
